#include <tinx/device.h>
#include <tinx/part.h>
#include <tinx/assert.h>
#include <tinx/arena.h>

typedef struct part_entry_t
{
    u8 pretty;
    u8 head_start;
    u8 sector_start : 6;
    u16 chs_start : 10;
    u8 type;        // 分区类型 (文件系统)
    u8 head_end;
    u8 sector_end : 6;
    u16 chs_end : 10;
    u32 lba;          // 分区起始 LBA 号
    u32 sectors; // 分区扇区数
}_packed part_entry_t;

typedef struct boot_sector_t
{
    u8 text[446];   // 代码段
    struct part_entry_t part_table[4];
    u16 magic;      // 魔数
}_packed boot_sector_t;

typedef struct part_t
{
    u32 sectors;
    u32 lba;
    dev_t dev;
    dev_t disk;
    int type;
} part_t;

static part_t part_table[256];

part_t *get_free_part()
{
    for (size_t i = 0; i < 256; i++)
    {
        part_t *part = &part_table[i];
        if (part->dev == EOF)
            return part;
    }
    panic("No more parts!");
}

int part_ioctl(part_t *part, int cmd, void *args, int flags)
{
    switch (cmd)
    {
    case DEV_CMD_SECTORS_COUNT:
        return part->sectors;
    case DEV_CMD_SECTORS_START:
        return part->lba;
    case DEV_CMD_PART_TYPE:
        return part->type;
    }
    return EOF;
}

// 事实上，因为磁盘调度算法只用读硬盘，这个函数并不会被调用
int part_read(part_t *part, char *buf, int count, u32 idx)
{
    if (idx + count >= part->sectors)
    {
        return EOF;
    }
    return device_read(part->dev, buf, count, part->lba + idx, 0);
}

// 事实上，这个函数并不会被调用
int part_write(part_t *part, char *buf, int count, u32 idx)
{
    if (idx + count >= part->sectors)
    {
        return EOF;
    }
    return device_write(part->dev, buf, count, part->lba + idx, 0);
}

void part_scan(dev_t dev, dev_t parent)
{
    boot_sector_t *boot = (boot_sector_t *)kmalloc(sizeof(boot_sector_t));
    device_read(dev, (char *)boot, 1, 0, 0); // 读取 MBR 所在的扇区
    u32 lba = device_ioctl(dev, DEV_CMD_SECTORS_START, 0);

    for (size_t idx = 0; idx < 4; idx++)
    {
        part_entry_t *entry = &boot->part_table[idx];

        if (entry->type == PART_FS_NULL)
            continue;

        part_t *part = get_free_part();
        part->disk = parent; 
        part->lba = entry->lba + lba;
        part->sectors = entry->sectors;
        part->type = entry->type;

        part->dev = device_create(
            DEV_TYPE_BLOCK, DEV_FLAG_PART, parent, NULL, part,
            part_write, part_read, part_ioctl);

        if (entry->type == PART_FS_EXTENDED)
        {
            part_scan(part->dev, parent);
        }
    }
    kfree(boot);
}

void part_init()
{
    for (size_t i = 0; i < 256; i++)
    {
        part_t *part = &part_table[i];
        part->dev = EOF;
    }
}